#
# Copyright (c) 2009, 2010, Davide Cittaro
# All rights reserved.
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of the <organization> nor the
#       names of its contributors may be used to endorse or promote products
#       derived from this software without specific prior written permission.
# 
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

# This module includes most of the helper function
# which I coulnd't put into classes...

import re, sys, os.path
from time import ctime
from config import Enzymes, Modifications, Elements, XLinkers
import sequence
import MSMS
import numpy as np
from hashlib import md5
import scipy.stats as SS

def timestring():
  """ A custom format for time, enclosed by brackets
  for debug
  """
  return '[ ' + ctime() + ' ]'

def Debug(nl= 1,*msg):
  """ A function to print to stderr the debug info
  Input can be assed as a list of arguments here casted to strings
  Terminator is defined by the first value passed
  1 -> \n
  2 -> \r
  3 -> \b
  """
  msg = ' '.join([str(x) for x in  msg])
  if nl == 1:
    msg += '\n'
  elif nl == 2:
    msg += '\r'
  elif nl == 3:
    msg += '\b'
  sys.stderr.write(msg)


def outputCoordinates(interaction):
  """ Prints coordinates for the given interaction
  """
  (prot1, prot2) = interaction.linkedProteins()
  outputString = prot1.accession() + ':' + str(interaction.xPositions()[0]) + ' ' + prot2.accession() + ':' + str(interaction.xPositions()[1])
  return outputString

def outputModifications(peptide, options):
  """ Prints modification string  for the given peptide
  """
  outputString = "[ "
  for pos, mod in peptide.modifications().items():
    if mod in options["fixedMods"]: continue
    outputString = outputString + mod + ':' + str(pos) + ','  
  outputString = outputString[:-1] + " ]"
  return outputString

def outputError(sp, xl):
  """ Prints the mass error in ppm
  """
#  return "%.2f" % ((sp.precursorMass() - xl.mass()) / sp.precursorMass() * 1000000)
  return "%.2f" % ((sp.precursorMass() - xl.mass()) / sp.precursorMass() * 1000000)

def outputParameters(dictOptions):
  """ Prints the operational parameters
  """
  print "# Spectra files:", ','.join([os.path.basename(x) for x in dictOptions["spectraFiles"]])
  print "# Sequence file:", os.path.basename(dictOptions["fastaFile"])
  print "# Enzyme:", dictOptions["enzyme"]
  print "# Missed cleavages:", dictOptions["missedCleavages"]
  print "# Min. peptide length:", dictOptions["minLength"]
  print "# Max. peptide length:", dictOptions["maxLength"]
  print "# MS Tolerance:", dictOptions["ms1tolerance"], dictOptions["unit1"]
  print "# MS/MS Tolerance:", dictOptions["ms2tolerance"], dictOptions["unit2"]
  print "# Fixed modifications:", ','.join(dictOptions["fixedMods"])
  print "# Variable modifications:", ','.join(dictOptions["varMods"])
  print "# Max. number of variable modifications:", dictOptions["maxVariableModifications"]
  print "# X-linkers:", ','.join(dictOptions["xlinker"])
  print "# Number of best spectra to retain:", dictOptions["retainBestSpectra"]
  print "# p-value filter:", dictOptions["pValueFilter"]
  print "# Match score filter:", dictOptions["mScoreFilter"]
  print "# Min. interaction length:", dictOptions["interactionLengthFilter"]
  print "# Number of top matching ions:", dictOptions["topIons"], "every", dictOptions["topWindow"], "Da"
  print "# Min. number of matching ions per peptide:", dictOptions["minMatchingIons"]
  print "# Percentile of MMS:", dictOptions["MMSFilter"]
  print "# Max. G Score:", dictOptions["GScoreFilter"]
  print "# Save numpy matrices:", dictOptions["saveNPData"]
  print "# Load numpy matrices:", dictOptions["loadNPData"]
  print

  
def buildModString(peptide):
  """ Returns the modification string for a peptide
  """
  retString = ''
  moddict = peptide.modifications()
  for (k, i) in peptide.modifications().items():
    retString +=':'.join([str(x) for x in [k, i]])
  return retString

def checkSumXlink(p1, p2, xl):
  """ Calculates the md5 checksum for the xlink given by p1 and p2.
  The order is not relevant here.
  """
  l = [p1.checksum(), p2.checksum()]
  l.sort()  # we take the xlink in alphabetical order, so that we always get one order...
  l.append(xl)
  return md5('_'.join(l)).hexdigest()
  
def shufflePeptides(Proteins, MSMS, options):
  """ This function creates the xlinks by shuffling and crossinkg
  all peptides defined for the proteins in input.
  
  Proteins: A list of Protein() objects
  MSMS: A MSMSAnalysis() object

  Returns a list of xlink() objects
  """
  # this function should be threaded

  MassList = MSMS.massList()
  XLinkList = np.array([], dtype=object)
  chklist = set()  #collects pep1/pep2 check for the homodimer to avoid duplicates (i.e. pep2/pep1)
  
#  Debug(3, "=")
  nTot = 0
  nValid = 0
  allPeps = np.array([], dtype=object)
  for plist in [x.peptides() for x in  Proteins]:
    allPeps = np.append(allPeps, plist)


  for xl in options["xlinker"]:
    # keep only peps that can xlink for this round
    retained = [pep for pep in allPeps if  pep.canXLink(xl)]
    nOfRetained = len(retained)
    for x in xrange(nOfRetained - 1):
      # loop over peptides on the triangular matrix
      pep1 = retained[x]
      for y in xrange(x + 1, nOfRetained):
        pep2 = retained[y]
        if pep1.checksum() == pep2.checksum(): 
          # we exclude dimers of the same peptide simply because
          # we cannot tell if it is a xlink or peptide with half charge...
          continue
        Debug(2, timestring(), "Creating xlinks [", nValid, "/", nTot, "]")
        nTot += 1  
        chksum = checkSumXlink(pep1, pep2, xl)  # get the md5 for this xlink (by pep sequencees and modifications)
        if chksum in chklist: continue      # if we have already seen this xlink, jump over
        chklist.add(chksum)
        mass = np.double(pep1.molecularMass() + pep2.molecularMass() + XLinkers.getfloat(xl, 'mmass'))
        if options["unit1"] == "ppm":
          maxErr = mass * options["ms1tolerance"] / 1000000.0 
        else:
          maxErr = options["ms1tolerance"]
        # get a mask of spectra that have mass compatible with this xlinks
        mask = (MassList >= mass - maxErr) & (MassList <= mass + maxErr)
        if not mask.any(): continue
        # Build a new xlink
        tempXLink = sequence.XLink()
        tempXLink.setLPeptide(pep1)
        tempXLink.setRPeptide(pep2)
        tempXLink.setMass(mass)
        tempXLink.setXlinker(xl)
        # right now check the best xlink position
        # if after that no ion match, discard the spectrum
        for xbool in xrange(len(mask)):
          if not mask[xbool]: continue # skip spectra with "wrong" mass
          sp = MSMS.spectra()[xbool]
          tempXLink.guessX(sp, options)
#              if tempXLink.numberOfMatchingIons(sp) < options["minMatchingIons"]:
           # Skip spectra if there are less than a certain number of ions
          if tempXLink.lPeptide().numberOfMatchingIons(sp) < options["minMatchingIons"]:
            mask[xbool] = False
          if tempXLink.rPeptide().numberOfMatchingIons(sp) < options["minMatchingIons"]:
            mask[xbool] = False
        # Use mask as array index to retain only good spectra
        tempXLink.setSpectra(MSMS.spectra()[mask])  # many spectra can explain this xlink...
        tempXLink.setCheckSum(chksum)
        if tempXLink.numberOfMatchingSpectra():
          XLinkList = np.append(XLinkList, tempXLink)
          nValid += 1
          [sp.addMatch(tempXLink) for sp in tempXLink.spectra()]

  del chklist     

  return XLinkList

def pepseqMap(protList):
  """ Returns a dictionary which connects each peptide to a protein
  and the position in it  
  """

  pep2prot = {}
  
  for protein in protList:
    protSeq = protein.sequence()
    for peptide in protein.peptides():
      pepSeq = peptide.sequence()
      posList = [x.start() + 1 for x in re.finditer(pepSeq, protSeq)]
      try:
        pep2prot[pepSeq][protein] = posList
      except KeyError:
        pep2prot[pepSeq] = {protein: posList}
  return pep2prot

def poissonMean(xl, sp, options):
  boundMasses = matchingSpectrum(xl, sp, options,exact = True).peakMasses()
  mean = 0.0
  
  for x in range(1, len(boundMasses)):
    l = boundMasses[x - 1]
    h = boundMasses[x]
    mean += sp.slice(l, h).topIons(options["topIons"],options["topWindow"]).length()
      
  mean /= len(boundMasses)
  
#  mean = (sp.peakMasses()[-1] - sp.peakMasses()[0])/options["topWindow"] * options["topIons"] / mean
  return mean
  

def buildInteractionList(xlList):
  """ Build the interactions from a list of XLink() object
  Returns the list of all interactions
  """
  
  xedict = {}
  for xl in xlList:
    for p1 in xl.lPeptide().positionsInProtein():
      for p2 in xl.rPeptide().positionsInProtein():
        c = xl.positionCheckSum(xl.lPeptide(), xl.rPeptide(), p1 + xl.lXPos(), p2 + xl.rXPos())
        try:
          xedict[c].addXLink(xl)
        except (TypeError, KeyError):
          xedict[c] = sequence.Interaction()
          xedict[c].addXLink(xl)
          xedict[c].setPositionCheckSum(c)
          xedict[c].setXPos1(p1 + xl.lXPos())
          xedict[c].setXPos2(p2 + xl.rXPos())
        Debug(2, timestring(), "Building interactions [", len(xedict), "]")
  
  return np.array(xedict.values())
      
def scoreInteractions(InteractionList, options):
  """ Calculates the G-Score for each interaction.

  G = 2 * \sum(i=1_r) k_i * ln(p_0/p_i)
  where
  r = number of xlinks
  k_i = number of evidence of xlink
  p_0 = null prob for xlink
  p_i = prob of xlink
  
  Also check if the interaction should be filtered out
  """
  
  MMS = []
  
  for interaction in InteractionList:
    MMS.append(interaction.aggregateMatchScore())
    G = 0.0
    for xl in interaction.xlinks():
      k_i = xl.numberOfMatchingSpectra()
      for sp in xl.spectra():
        p_0 = 1.0 / len(sp.matches())
        p_i = xl.pValue(sp)
        G = G + k_i * np.log(p_i / p_0)
    interaction.setGScore(2 * G)
    if interaction.GScore() > options["GScoreFilter"]:
      interaction.filterOut()
  
  MMS_thr = SS.scoreatpercentile(MMS, options["MMSFilter"])
  [x.filterOut() for x in InteractionList if x.aggregateMatchScore() < MMS_thr]

def matchingSpectrum(xlink, spectrum, options, exact = True):
  """ return a spectrum only containing matching ions
  it can be exact, only containing matches, or not that means
  also surrounding peaks are included
  """
  masses = set()
  
  spOut = MSMS.ms2spectrum()
  spOut.setPrecursorMz(spectrum.precursorMz())
  spOut.setCharge(spectrum.charge())
  spOut.setDescription(spectrum.description())
  
  #b-ions
  for pep in [xlink.lPeptide(), xlink.rPeptide()]:
    brep = [int(x) for x in np.binary_repr(pep.matchedBIons(spectrum))]
    brep.reverse()
    for i, x in enumerate(brep):
      if x: masses.add(pep.bMass(i))
  #y-ions
    brep = [int(x) for x in np.binary_repr(pep.matchedYIons(spectrum))]
    brep.reverse()
    for i, x in enumerate(brep):
      if x: masses.add(pep.yMass(i))
      
  
  mlist = list(masses)
  mlist.sort()
 
  masses = set()
  if exact:
    # this if we want the peaks of the matching ions
    for m in mlist:
      cp = spectrum.closestTo(m, options["ms2tolerance"], options["unit2"])
      if not len(cp) or cp[0] in masses:
        # we already have added this mass, dont' add the peak
        continue
      masses.add(cp[0])
      spOut.addPeak(cp)
  else:
    # this if we want the peaks of the matching ions
    # and the spectra within tolerance...
    for m in mlist:
      # first get the masses of the ions
      # within the tolerance
      for im in spectrum.closeTo(m, options["ms2tolerance"], options["unit2"]).peakMasses():
        masses.add(im)
    mask = np.array([x in masses for x in spectrum.peakMasses()])    
    # now the mask contains a True if the ion in the original spectrum
    # matches somewhere, we can use it to filter the orignal masses
    # and the corresponding intensities
    spOut.setSpectrum([spectrum.peakMasses()[mask], spectrum.peakIntensities()[mask]], axis = 0)
  
  masses.clear()
  return spOut

def mad(data):
  # returns median absolute deviation
  data = data[np.isnan(data) ^ True]
  median = np.median(data)
  absmed = np.abs(data - median)
  return np.median(absmed)
  
def calcPValues(XList):
  """ Calculates the p-values for each xlink
  Intensity scores are used to build a lognormal distribution
  which is then used to get p-values
  """
  mscores = np.array([])
  for xl in XList:
    for sp in xl.spectra():
#      mscores = np.append(mscores, xl.matchScore(sp))
# try logs... apparently those distribute normally
      mscores = np.append(mscores, np.log(xl.matchScore(sp)))


  deviation = mad(mscores) #calculate the deviation as MAD
  cValue = np.median(mscores) # calculate the central value as median
  distr = SS.distributions.norm(cValue, deviation)

  for xl in XList:
    for sp in xl.spectra():
#      pvalue = 1 - distr.cdf(xl.matchScore(sp))
      pvalue = 1 - distr.cdf(np.log(xl.matchScore(sp)))
      xl.setPValue(pvalue, sp)
    xl.setWeight(cValue)

###
def printMatchingIons(xlink, spectrum, options):
# Copied and adapted from old silk implementation :-)
  tolerance = options["ms2tolerance"]
  unit = options["unit2"]
  lPeptide = xlink.lPeptide()
  rPeptide = xlink.rPeptide()
  
  NSpool = []
  CSpool = []
  
#  print lPeptide, xlink.xStart(), lPeptide.sequence(), lPeptide.sequence()[xlink.xStart()]
#  print rPeptide, xlink.xEnd(), rPeptide.sequence(), rPeptide.sequence()[xlink.xEnd()]
  #sp.closeTo(mz, tol, u).peakMasses()
  # for now print only b and y ions of simple peptides

  #Left Peptide
  brep = [int(x) for x in np.binary_repr(lPeptide.matchedBIons(spectrum))]
  brep.reverse()
  yrep = [int(x) for x in np.binary_repr(lPeptide.matchedYIons(spectrum))]
  yrep.reverse()
  
  for s in range(xlink.lXPos() - 1):
    index = s + 1
    mz = lPeptide.bMass(index)
    found = spectrum.closeTo(mz, tolerance, unit).closestTo(mz)
    if len(found):
      mString = "[ " + "%.3f" % found[0] + ", " + "%.3f" % found[1] + " ]"
    else:
      mString = "[ ]"
    ionMass = "%.3f" % lPeptide.bMass(index)
    row = 'b'+str(index) + ' ' +  lPeptide.sequence()[s] + ' ' + ionMass + ' ' +  mString
    NSpool.append(row)
  for s in range(xlink.lXPos(), lPeptide.length()):
    index = lPeptide.length() - s
    mz = lPeptide.yMass(index)
    found = spectrum.closeTo(mz, tolerance, unit).closestTo(mz)
    if len(found):
      mString = "[ " + "%.3f" % found[0] + ", " + "%.3f" % found[1] + " ]"
    else:
      mString = "[ ]"
    ionMass = "%.3f" % lPeptide.yMass(index)
    row = 'y'+str(index) + ' ' +  lPeptide.sequence()[s] + ' ' +  ionMass + ' ' + mString
    NSpool.append(row)

  for s in range(xlink.rXPos() - 1):
    index = s + 1
    mz = rPeptide.bMass(index)
    found = spectrum.closeTo(mz, tolerance, unit).closestTo(mz)
    if len(found):
      mString = "[ " + "%.3f" % found[0] + ", " + "%.3f" % found[1] + " ]"
    else:
      mString = "[ ]"
    ionMass = "%.3f" % rPeptide.bMass(index)
    row = 'b'+str(index) + ' ' + rPeptide.sequence()[s] + ' ' + ionMass + ' ' + mString
    CSpool.append(row)
  for s in range(xlink.rXPos(), rPeptide.length()):
    index = rPeptide.length() - s
    mz = rPeptide.yMass(index)
    found = spectrum.closeTo(mz, tolerance, unit).closestTo(mz)
    if len(found):
      mString = "[ " + "%.3f" % found[0] + ", " + "%.3f" % found[1] + " ]"
    else:
      mString = "[ ]"
    ionMass = "%.3f" % rPeptide.yMass(index)
    row = 'y'+str(index) + ' ' + rPeptide.sequence()[s] + ' ' + ionMass + ' ' + mString
    CSpool.append(row)
  
  NS = max([len(x) for x in NSpool])

  print lPeptide.sequence(), ' ' * (NS - lPeptide.length()), rPeptide.sequence()
#  print map(lambda x: np.binary_repr(x), lPeptide.matchedIons(spectrum)), map(lambda x: np.binary_repr(x), rPeptide.matchedIons(spectrum))
  for i in range(max(len(CSpool), len(NSpool))):
    
    try:
      l = len(NSpool[i])
      spaces = ' ' * (NS - l)
      try:
        print NSpool[i], spaces, CSpool[i]
      except IndexError:
        print
    except IndexError:
      spaces = ' ' * (1 + NS)
      print spaces, CSpool[i]
  


def matchIons(pep, sp, options, splitPos = -1):
  """ Given a peptide and a spectrum gives the
  matching ions as binary string and also
  calculate the intensity score based
  on negative binomial distribution of intensity values
  """
  matchScore = 0.0
  mScore = []
  IntPS = np.array([])	# an array with percentile value for intensities
  mcount = 0
  bMatch = 0
  yMatch = 0
  k = 0     # number of failures for NB: it is at least the length of peptide, otherwise the number of closest ions
  r = 0     # the number of successes: number of matching ions
  
  if sp.length() == 0:
    return (0.0, 0, 0)
  
  Intensities = sp.peakIntensities()
  
  if splitPos == -1:
    bSplit = pep.length()
    ySplit = 0
  else:
    bSplit = splitPos
    ySplit = pep.length() - splitPos + 1
  for idx in range(1 + pep.length()):
    revidx = pep.length() - idx
    if idx < bSplit:
      clsp = sp.closeTo(pep.bMass(idx), options["ms2tolerance"], options["unit2"])
      l = clsp.length()
      k += l
      if l:
        # there are matching ions in for this
        bMatch = bMatch + 2**idx
        # get the percentile for this intensity. We'll use it later to weight probabilities
        IR = SS.percentileofscore(Intensities, clsp.mostIntense()[1])
        IntPS = np.append(IntPS, IR)
        r += 1

    if revidx < ySplit:
      clsp = sp.closeTo(pep.yMass(revidx), options["ms2tolerance"], options["unit2"])
      l = clsp.length()
      k += l
      if l:
        yMatch = yMatch + 2**revidx
        IR = SS.percentileofscore(Intensities, clsp.mostIntense()[1])
        IntPS = np.append(IntPS, IR)
        r += 1

  if k < pep.length():
    k = pep.length()
  p = float(r) / k	# for negative binomial, this is the MLE of probability
  pmf = SS.nbinom.pmf(1, k - r, p)
  for ps in IntPS:
    # the negative binomial probabiliy is calculated for each ion, (k - r) is the number
    # of failures.
    matchScore += pmf * ps	# calculate the prob for each ion and weight by percentile

  if np.isnan(matchScore):
    # just in case...
    matchScore = 0.0
  return (matchScore, bMatch, yMatch)


def sillyScore(matchString):
  return max([len(x) for x in matchString.split('0')]) * sum([int(x) for x in matchString])

